import Color
import cv2
import numpy as np

"""
color = {'red_high': [255, 255, 255], 'red_low': [0, 0, 0],
         'green_high': [255, 255, 255], 'green_low': [0, 0, 0],
         'blue_high': [255, 255, 255], 'blue_low': [0, 0, 0]}
"""
# 读取颜色阈值
color = Color.color
color_list = ['red', 'green', 'blue']
color_value = {'red': [0, 0, 255], 'green': [0, 255, 0], 'blue': [255, 0, 0]}


# 更新颜色阈值
def Update_Color(colors):
    with open('./Color.py', 'w') as f:
        temp_s = "color = {}".format(str(colors))
        print(temp_s)
        f.write(temp_s)


def nothing(x):
    pass


# 创建滑动条
cv2.namedWindow('get_color')
cv2.createTrackbar('low_L', 'get_color', 0, 255, nothing)
cv2.createTrackbar('low_A', 'get_color', 0, 255, nothing)
cv2.createTrackbar('low_B', 'get_color', 0, 255, nothing)
cv2.createTrackbar('high_L', 'get_color', 255, 255, nothing)
cv2.createTrackbar('high_A', 'get_color', 255, 255, nothing)
cv2.createTrackbar('high_B', 'get_color', 255, 255, nothing)


# 获取滑动条阈值
def get_Trackbar_color():
    low_L = cv2.getTrackbarPos('low_L', 'get_color')
    low_A = cv2.getTrackbarPos('low_A', 'get_color')
    low_B = cv2.getTrackbarPos('low_B', 'get_color')
    high_L = cv2.getTrackbarPos('high_L', 'get_color')
    high_A = cv2.getTrackbarPos('high_A', 'get_color')
    high_B = cv2.getTrackbarPos('high_B', 'get_color')
    return [low_L, low_A, low_B], [high_L, high_A, high_B]


def get_color(cap):
    # 判断摄像头是否初始化并成功打开
    if cap.isOpened():
        state, frame = cap.read()  # 抓取下一个视频帧，返回状态以及图像
        while state:
            state, frame = cap.read()  # 抓取每一帧图像
            blur_image = cv2.blur(frame, (3, 3))  # 平滑图像
            lab_image = cv2.cvtColor(blur_image, cv2.COLOR_BGR2LAB)  # 转换颜色空间
            colorLow, colorHigh = get_Trackbar_color()  # 获取滑动条颜色阈值
            lab_mask = cv2.inRange(lab_image, np.array(colorLow), np.array(colorHigh))  # 颜色阈值分割
            cv2.imshow('image', lab_mask)  # 显示图像帧
            k = cv2.waitKey(25)  # 每25ms读取一次键盘按键，同时控制帧率
            if k == ord('q'):  # 当检测到键盘q键按下，程序退出
                cv2.destroyWindow('image')  # 销毁窗口
                break
            elif k == ord('r'):  # 修改颜色阈值
                color['red_high'] = colorHigh
                color['red_low'] = colorLow
            elif k == ord('g'):  # 修改颜色阈值
                color['green_high'] = colorHigh
                color['green_low'] = colorLow
            elif k == ord('b'):  # 修改颜色阈值
                color['blue_high'] = colorHigh
                color['blue_low'] = colorLow
            elif k == ord('s'):  # 更新颜色阈值
                Update_Color(color)


def color_detection(cap):
    # 判断摄像头是否初始化并成功打开
    if cap.isOpened():
        state, frame = cap.read()  # 抓取下一个视频帧，返回状态以及图像
        while state:
            state, frame = cap.read()  # 抓取每一帧图像
            lab_image = cv2.cvtColor(frame, cv2.COLOR_BGR2LAB)  # 颜色空间转换
            blur_image = cv2.blur(lab_image, (3, 3))            # 平滑滤波操作
            drawing = frame.copy()
            for col in color_list:
                # 颜色阈值分割
                lab_mask = cv2.inRange(blur_image, np.array(color[col + '_low']), np.array(color[col + '_high']))
                Kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
                lab_mask = cv2.morphologyEx(lab_mask, cv2.MORPH_OPEN, Kernel)   # 开操作，消除较小的白色噪点
                # 查找轮廓
                contours, hierarchy = cv2.findContours(lab_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)[-2:]
                # 找最大面积轮廓
                area_max = 0
                _cnt = None
                if len(contours) != 0:
                    for cnt in contours:
                        area = cv2.contourArea(cnt)
                        if area > area_max:
                            area_max = area
                            _cnt = cnt
                    drawing = cv2.drawContours(drawing, [_cnt], -1, color_value[col], 2, 8)     # 绘制轮廓
                    rect = cv2.minAreaRect(_cnt)  # 得到最小外接矩形
                    box = np.int0(cv2.boxPoints(rect))
                    cv2.drawContours(drawing, [box], 0, color_value[col], 2)    # 绘制矩形
                cv2.imshow('image', drawing)  # 显示图像帧
            k = cv2.waitKey(25)  # 每25ms读取一次键盘按键，同时控制帧率
            if k == ord('q'):  # 当检测到键盘q键按下，程序退出
                break
            elif k == ord('c'):
                get_color(cap)  # 更新颜色阈值


if __name__ == "__main__":
    # 获取摄像头对象
    cap = cv2.VideoCapture(0)
    # 颜色识别
    color_detection(cap)
